# coding: utf-8
# Project：PyCharmMiscProject
# File：实验三.Socket编程.py
# Author：liuchuanzhi
# Date ：2025/4/16 18:33
# IDE：PyCharm
import socket
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import os

# 创建 socket 对象
client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

# 连接到服务器
server_address = ('192.168.197.42',8888)
client_socket.connect(server_address)

# 接收服务器发送的密钥和IV
key_iv = client_socket.recv(32)
key = key_iv[:16]
iv = key_iv[16:]


def encrypt_message(message):
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted = cipher.encrypt(pad(message.encode('utf-8'), AES.block_size))
    return iv + encrypted


def decrypt_message(encrypted_message):
    iv = encrypted_message[:AES.block_size]
    encrypted = encrypted_message[AES.block_size:]
    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted = unpad(cipher.decrypt(encrypted), AES.block_size)
    return decrypted.decode('utf-8')


try:
    conversation = []
    while True:
        # 发送消息给服务器
        plaintext_message = input('请输入要发送给服务器的消息: ')
        encrypted_message = encrypt_message(plaintext_message)
        client_socket.send(encrypted_message)
        conversation.append(f"客户端: {plaintext_message}")

        # 接收服务器加密消息
        encrypted_data = client_socket.recv(1024)
        if not encrypted_data:
            break
        # 解密消息
        decrypted_message = decrypt_message(encrypted_data)
        conversation.append(f"服务端: {decrypted_message}")
        print('服务器加密消息解密后: {}'.format(decrypted_message))

        if input("是否继续对话？Y/N").upper() == "N":
            break

    # 将对话保存为加密文件
    conversation_text = '\n'.join(conversation)
    encrypted_conversation = encrypt_message(conversation_text)
    with open('encrypted_conversation.bin', 'wb') as file:
        file.write(encrypted_conversation)
    print("对话已保存为加密文件 encrypted_conversation.bin")

    # 后续操作可以继续添加，这里仅打印对话记录
    print("对话记录：")
    for line in conversation:
        print(line)

except Exception as e:
    print(f'发生错误: {e}')
finally:
    # 关闭连接
    client_socket.close()

